from collections import defaultdict, OrderedDict
import numpy as np
import sys

class ModSeminario:
    def __init__(self, logfile, fchkfile, vibrational_scaling=1):
        self.logfile = logfile
        self.fchkfile = fchkfile
        self.vibrational_scaling_squared = vibrational_scaling**2
        self.ele2num = {'H':1,'He':2,'Li':3,'Be':4,'B':5,'C':6,'N':7,'O':8,'F':9,'Ne':10,
                        'Na':11,'Mg':12,'Al':13,'Si':14,'P':15,'S':16,'Cl':17,'Ar':18,
                        'K':19,'Ca':20,'Ti':22,'V':23,'Cr':24,'Mn':25,'Fe':26,'Co':27,
                        'Ni':28,'Cu':29,'Zn':30,'Ga':31,'Ge':32,'As':33,'Se':34,'Br':35,
                        'Rh':45,'Pd':46,'Pt':78,'Ag':47,'Au':79,'Cd':48,'Hg':80,'I':53,
                        'Ir':77,'Pb':82,'Ce':58}
        self.num2ele = dict((value,key) for key,value in self.ele2num.items())
        self.modified_Seminario_method()

    def coords_from_fchk(self):
        '''Function extracts xyz file from the .fchk output file from Gaussian, this
        provides the coordinates of the molecules'''
        fchk = open(self.fchkfile, 'r')
        numbers = []
        read_numbers = 0
        list_coords = []
        read_coords = 0
        hessian = []
        read_hessian = 0
        #Get atomic number from fchk
        for line in fchk:
            if 'Atomic numbers' in line:
                read_numbers = 1
                continue
            if read_numbers == 1:
                if line.strip()[:1].isalpha():
                    read_numbers = 0
                else:
                    numbers = numbers + line.split()
            #Gets coordinates
            if 'Current cartesian coordinates' in line:
                read_coords = 1
                continue
            if read_coords == 1:
                if line.strip()[:1].isalpha():
                    read_coords = 0
                else:
                    list_coords = list_coords + [float(i) for i in line.split()]
            #Gets Hessian
            if 'Cartesian Force Constants' in line:
                read_hessian = 1
                continue
            if read_hessian == 1:
                if line.strip()[:1].isalpha():
                    read_hessian = 0
                else:
                    hessian = hessian + [float(i) for i in line.split()]
        list_coords = np.array(list_coords) * 0.529
        natom = len(list_coords) / 3
        coords = np.reshape(list_coords, (-1, 3))
        elements = [[self.num2ele[int(i)]] for i in numbers]
        self.coords = coords
        self.elements = elements
        self.hessian = hessian
        if not hessian:
            print('Error! Hessian data not extracted from fchk file {:s}\n'
                  'You need to make sure its a freq job\n'
                  'and check "Cartesian Force Constants" in fchk file'
                  .format(self.fchkfile))
            sys.exit()
        self.natom = int(natom)
        #Get coords from fchk

    def bond_angle_lists(self):
        '''This function extracts a list of bond and angles from the Gaussian .log file'''
        log = open(self.logfile,'r')
        read_ba = 0
        for line in log:
            if 'Name  Definition              Value          Derivative Info.' in line:
                read_ba = 1
                bondlist = []
                anglelist = []
                dihedrallist = []
                continue
            if read_ba == 1 and '-----' in line:
                read_ba = 2
                continue
            if read_ba == 2 and '-----' in line:
                read_ba = 0
                continue
            if read_ba == 2:
                le = line.strip().split()
                if le[1].startswith("R"):
                    bondlist.append([int(i) for i in le[2].strip('R()').split(',')])
                if le[1].startswith("A") and le[2].startswith("A"):
                    anglelist.append([int(i) for i in le[2].strip('A()').split(',')])
                if le[1].startswith("D"):
                    dihedrallist.append(le[2].strip('D()').split(','))
        self.bondlist = bondlist
        self.anglelist = anglelist

    def input_data_processing(self):
        '''This function takes input data that is need from the files supplied
        Function extracts input coords and hessian from .fchk file, bond and angle
        lists from .log file and atom names if a z-matrix is supplied
        Gets Hessian in unprocessed form and writes .xyz file too '''

        self.coords_from_fchk()
        self.bond_angle_lists()
        hessian = np.zeros((3*self.natom,3*self.natom))
        # Change from Hartree/bohr to kcal/mol /ang
        unprocessed_hessian = np.array(self.hessian) * 627.509391 / (0.52917721067121**2)
        # origin algorithm
        # for i in range(3*self.natom):
        #     print(hessian[i][0:i+1])
        #     a = int((i+1)*i/2)
        #     b = int((i+1)*(i+2)/2)
        #     hessian[i][0:i+1] = unprocessed_hessian[a:b]
        # Change from 1d array to low triangle matrix
        hessian[np.tril_indices(3*self.natom)] = unprocessed_hessian
        # change from low triangle matrix to symmetrical matrix
        self.hessian = np.where(hessian,hessian,hessian.T)

    def force_constant_bond(self, a, b):
        '''#Force Constant - Equation 10 of Seminario paper - gives force
           constant for bond
           a and b are atom index
           '''
        eigenvalues_AB = self.eigenvalues[a, b, :]
        eigenvectors_AB = self.eigenvectors[0:3,0:3, a, b]
        diff_AB = self.coords[a] - self.coords[b]
        unit_vectors_AB = diff_AB / np.linalg.norm(diff_AB)
        k_AB = 0
        for i in range(3):
            k_AB = k_AB + eigenvalues_AB[i] * abs(np.dot(unit_vectors_AB,eigenvectors_AB[:,i]))
        k_AB = -1 * k_AB/2.0
        return k_AB

    def bonds_calculated_printed(self):
        '''This function uses the Seminario method to find the
        bond parameters and print them to file'''
        bond_s2p = OrderedDict()
        bond_length_list = np.zeros((1,len(self.bondlist)))
        for bond in self.bondlist:
            a = bond[0] - 1
            b = bond[1] - 1
            AB = self.force_constant_bond(a, b)
            BA = self.force_constant_bond(b, a)
            # k_b = np.real((AB+BA)/2*self.vibrational_scaling_squared)*4.184*100*2 #for gmx
            k_b = np.real((AB+BA)/2*self.vibrational_scaling_squared)
            # r_b = np.linalg.norm(self.coords[a]-self.coords[b])/10 # for gmx
            r_b = np.linalg.norm(self.coords[a]-self.coords[b])
            bond_s2p[(str(bond[0]),str(bond[1]))] = (r_b,k_b)
        self.bond_s2p = bond_s2p

    def u_PA_from_angles(self,atom_A,atom_B,atom_C):
        # This gives the vector in the plane A,B,C and perpendicular to A to B
        # A B C are atom  serial.

        diff_AB = self.coords[atom_B-1] - self.coords[atom_A-1]
        u_AB = diff_AB / np.linalg.norm(diff_AB)
        diff_CB = self.coords[atom_B-1] - self.coords[atom_C-1]
        u_CB = diff_CB / np.linalg.norm(diff_CB)
        cross = np.cross(u_CB, u_AB)
        u_N = cross / np.linalg.norm(cross)
        u_PA = np.cross(u_N, u_AB)
        u_PA = u_PA/ np.linalg.norm(u_PA)
        return u_PA

    def force_angle_constant(self,atom_A,atom_B,atom_C,scale_factor):
        '''Force Constant- Equation 14 of seminario calculation paper - gives force
        constant for angle (in kcal/mol/rad^2) and equilibrium angle in degrees'''
        a = atom_A - 1
        b = atom_B - 1
        c = atom_C - 1
        diff_AB = self.coords[b] - self.coords[a]
        u_AB = diff_AB / np.linalg.norm(diff_AB)
        diff_CB = self.coords[b] - self.coords[c]
        u_CB = diff_CB / np.linalg.norm(diff_CB)
        bond_length_AB = self.bond_lengths[a][b]
        eigenvalues_AB = self.eigenvalues[a, b, :]
        eigenvectors_AB = self.eigenvectors[0:3, 0:3, a, b]
        bond_length_CB = self.bond_lengths[b][c]
        eigenvalues_CB = self.eigenvalues[c, b, :]
        eigenvectors_CB = self.eigenvectors[0:3, 0:3, c, b]
        # Normal vector to angle plane found
        u_N = np.cross(u_CB, u_AB)
        u_N = u_N / np.linalg.norm(u_N)
        u_PA = np.cross(u_N, u_AB)
        u_PA = u_PA / np.linalg.norm(u_PA)
        u_PC = np.cross(u_N, u_CB)
        u_PC = u_PC / np.linalg.norm(u_PC)
        # Projections of eigenvalues to the norm of bond AB and bond CB
        sum_first = eigenvalues_AB * np.abs(np.dot(u_PA,eigenvectors_AB))
        sum_second = eigenvalues_CB * np.abs(np.dot(u_PC,eigenvectors_CB))
        # sum and Scaling due to additional angles - Modified Seminario Part
        sum_first = np.sum(sum_first)/scale_factor[(atom_B,atom_A,atom_C)]
        sum_second = np.sum(sum_second)/scale_factor[(atom_B,atom_C,atom_A)]
        # Added as two springs in series
        k_theta =  (1/(bond_length_AB**2*sum_first) + 1/(bond_length_CB**2*sum_second))**-1
        k_theta = abs(k_theta*0.5)
        # Equilibrium Angle
        theta_0 = np.rad2deg(np.arccos(np.dot(u_AB, u_CB)))
        return k_theta,theta_0

    def angles_calculated_printed(self):
        '''This function uses the modified Seminario method to find the angle'''
        angle_s2p = OrderedDict()

        # A structure is created with the index giving the central atom of the
        # angle, an array then lists the angles with that central atom.
        # ie. central_atoms_angles{3} contains an array of angles with central atom 3
        central_atoms_angles = OrderedDict()
        for i,a in enumerate(self.anglelist):
            angle_list = [[a[0],a[2],i],[a[2],a[0],i]]
            if a[1] in central_atoms_angles:
                central_atoms_angles[a[1]] += angle_list
            else:
                central_atoms_angles[a[1]] = angle_list
        # For the angle at central_atoms_angles[i] the corresponding
        # u_PA is the vector in ABC plane and perpendicular to AB, where ABC
        # corresponds to the order of the arguements
        # This is why the reverse order was also added'''
        unit_PA_all_angles = OrderedDict()
        for k,v in central_atoms_angles.items():
            unit_PA_all_angles[k] = []
            for l in v:
                unit_PA_all_angles[k].append(self.u_PA_from_angles(l[0],k,l[1]))
        # Finds the contributing factors from the other angle terms
        # Goes through the list of angles with the same central atom
        # And computes the term need for the modified Seminario method
        # if two angle in same plane share a same bond, then the scaling factor for this bond is 2
        # if two angle in perpendicular plane share a same bond, then the scaling factor for this bond is 1
        scaling_factor_all_angle = OrderedDict()
        for k,v in central_atoms_angles.items(): # k is central atom serial, v this [atom_a atom_c index_in_anglelist] for k
            for i1,l1 in enumerate(v):
                additional_contributions = []
                for i2, l2 in enumerate(v):
                    if i1 != i2 and l1[0] == l2[0]:
                        additional_contributions.append(abs(np.dot(unit_PA_all_angles[k][i1],unit_PA_all_angles[k][i2])))
                if len(additional_contributions) > 0:
                    scaling_factor = 1 + np.mean(additional_contributions)
                else:
                    scaling_factor = 1
                if (k,l1[0],l1[1]) not in scaling_factor_all_angle:
                    scaling_factor_all_angle[(k,l1[0],l1[1])] = scaling_factor
                else:
                    print("sth wrong")
        #
        # Finds the angle force constants with the scaling factors included for each angle
        self.angle_s2p = OrderedDict()
        for a in self.anglelist:
            ABC_k, ABC_a = self.force_angle_constant(a[0],a[1],a[2],scaling_factor_all_angle)
            CBA_k, CBA_a = self.force_angle_constant(a[2],a[1],a[0],scaling_factor_all_angle)
            # k_theta = (ABC_k+CBA_k) * 4.184  # gmx unit
            k_theta = (ABC_k+CBA_k) / 2
            self.angle_s2p[tuple([str(i) for i in a])] = (ABC_a,k_theta)

    def modified_Seminario_method(self):
        '''Program to implement the Modified Seminario Method
           origin matlab version Written by Alice E. A. Allen, TCM, University of Cambridge
           Reference using AEA Allen, MC Payne, DJ Cole, J. Chem. Theory Comput. (2018), doi:10.1021/acs.jctc.7b00785
           python version Written by Zhong Cheng, Wuhan university'''
        self.input_data_processing()
        # Find bond lengths
        bond_lengths = np.zeros((self.natom,self.natom))
        for i in range(self.natom):
            for j in range(self.natom):
                bond_lengths[i][j] = np.linalg.norm(self.coords[i]-self.coords[j])
        # Eigenvectors and eigenvalues calculated
        eigenvectors = np.zeros((3, 3, self.natom, self.natom),dtype = 'complex_')
        eigenvalues = np.zeros((self.natom,self.natom,3),dtype = 'complex_')
        for i in range(self.natom):
            for j in range(self.natom):
                A,B = np.linalg.eig(self.hessian[i*3:i*3+3,j*3:j*3+3])
                eigenvalues[i,j,:] = A
                eigenvectors[:,:,i,j] = B
        self.bond_lengths = bond_lengths
        self.eigenvalues = eigenvalues
        self.eigenvectors = eigenvectors
        self.bonds_calculated_printed() # generate self.bond_s2p
        self.angles_calculated_printed() #generate self.angle_s2p
        self.params = defaultdict(str)
        self.params['bonds'] = self.bond_s2p
        self.params['angles'] = self.angle_s2p
        # self.params['coords'] = self.coords
